FastAPIのWebSocket接続にサブプロトコルを使用したい

荒木智仁

2023.01.06

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

データアナリティクス事業本部の荒木です。

FastAPIでサブプロトコルを使用してWebSocket接続する機会があり、あまり該当する情報が無く大変だったのでまとめたいと思います。

環境

Windows
WSL2(Ubuntu 20.04.4 LTS)
Python(3.8.10)

本題

WebSocket接続の検証のため、クライアント側を簡単なJavascriptを作成して行いました。サーバ側とクライアント側のソースコードは以下になります。

サーバー

from fastapi import APIRouter, WebSocket

router = APIRouter()

@router.websocket(
    "/api/subscribe/"
)
async def websocket_endpoint(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            data = await ws.receive_text()
            await ws.send_text(f"data: {data}")
    except:
        await ws.close()

クライアント

<!DOCTYPE html>
<html>
    <head>
        <title>Chat</title>
    </head>
    <body>
        <h1>WebSocket Chat</h1>
        <ul id='messages'>
        </ul>
        <script>
            var ws = new WebSocket("ws://localhost:8000/api/subscribe/);
            ws.onmessage = function(event) {
                var messages = document.getElementById('messages')
                var message = document.createElement('li')
                var content = document.createTextNode(event.data)
                message.appendChild(content)
                messages.appendChild(message)
            };
            function sendMessage(event) {
                var input = document.getElementById("messageText")
                ws.send(input.value)
                input.value = ''
                event.preventDefault()
            }
        </script>
    </body>
</html>

python実行環境にてuvicorn main:app --reloadのコマンドでアプリケーションを実行します。
クライアントソースをブラウザにて開きます。
ボックスに入力した文字がサーバから返ってくることが確認できます。

サブプロトコルを使用するためクライアント側を以下のように変更します。

クライアント

<!DOCTYPE html>
<html>
    <head>
        <title>Chat</title>
    </head>
    <body>
        <h1>WebSocket Chat</h1>
        <ul id='messages'>
        </ul>
        <script>
            var ws = new WebSocket("ws://localhost:8000/api/subscribe/, [token1, token2]);
            ws.onmessage = function(event) {
                var messages = document.getElementById('messages')
                var message = document.createElement('li')
                var content = document.createTextNode(event.data)
                message.appendChild(content)
                messages.appendChild(message)
            };
            function sendMessage(event) {
                var input = document.getElementById("messageText")
                ws.send(input.value)
                input.value = ''
                event.preventDefault()
            }
        </script>
    </body>
</html>

上記で同様にアプリケーションを実行した際に、リクエストのどこにサブプロトコル情報が含まれるかわからず苦労しました。。。
結論、サブプロトコル情報はスコープの中にsubprotocolsとしてリスト形式で含まれます。

{'type': 'websocket', 'asgi': {'version': '3.0', 'spec_version': '2.3'}, 'http_version': '1.1', 'scheme': 'ws', 'server': ('127.0.0.1', 8000), 'client': ('127.0.0.1', 42678), 'root_path': '', 'path': '/api/subscribe/', 'raw_path': b'/api/subscribe/', 'query_string': b'', 'headers': [(b'host', b'localhost:8000'), (b'connection', b'Upgrade'), (b'pragma', b'no-cache'), (b'cache-control', b'no-cache'), (b'user-agent', b'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/108.0.0.0 Safari/537.36'), (b'upgrade', b'websocket'), (b'origin', b'null'), (b'sec-websocket-version', b'13'), (b'accept-encoding', b'gzip, deflate, br'), (b'accept-language', b'ja,en-US;q=0.9,en;q=0.8'), (b'sec-websocket-key', b'CwUA3B6CE2Y5WHZwHX12Iw=='), (b'sec-websocket-extensions', b'permessage-deflate; client_max_window_bits'), (b'sec-websocket-protocol', b'token1, token2')], 'subprotocols': ['token1', 'token2'], 'app': <fastapi.applications.FastAPI object at 0x7f68cb825280>, 'fastapi_astack': <contextlib.AsyncExitStack object at 0x7f68c2a8e880>, 'router': <fastapi.routing.APIRouter object at 0x7f68cb13eb20>, 'endpoint': <function websocket_endpoint at 0x7f68cb12cdc0>, 'path_params': {}, 'route': APIWebSocketRoute(path='/api/subscribe/', name='websocket_endpoint')}

そのためサーバー側ではサブプロトコルの情報をレスポンスで返すよう変更する必要があります。

from fastapi import APIRouter, WebSocket

router = APIRouter()

@router.websocket(
    "/api/subscribe/"
)
async def websocket_endpoint(ws: WebSocket):
    token = ws.scope["subprotocols"]
    for i, name in enumerate(token):
        if name == "token1":
            await ws.accept(token[i])
            try:
                while True:
                    data = await ws.receive_text()
                    await ws.send_text(f"data: {data}")
            except:
                await ws.close()
    else:
        await ws.close()

上記のように特定のプロトコルが含まれている場合、処理を実行するような処理にすることも可能です。